#include "kernel_operator.h"
#include "op_common.h"
namespace AscendC {

template <typename _DT_GRAD, typename _DT_X, typename _DT_TANHX, typename _DT_Y>
class Kernel {
public:
  TPipe pipe;
  int64_t size = 0;
  int64_t tileLength = 0;
  int64_t finalLength = 0;
  bool hasTanhx = false;
  DefInTensor(GRAD);
  DefInTensor(X);
  DefInTensor(TANHX);
  DefOutTensor(Y);

  DefBufVECCALC(CASTED_GRAD);
  DefBufVECCALC(CASTED_X);
  DefBufVECCALC(CASTED_TANHX);
  DefBufVECCALC(CASTED_Y);

  DefBufVECCALC(TEMP);
  DefBufVECCALC(TEMP2);
  DefBufVECCALC(TEMP3);

public:
  __aicore__ inline Kernel() {}

  template <typename T>
  __aicore__ inline void Init(GM_ADDR grad, GM_ADDR x, GM_ADDR tanhx, GM_ADDR y,
                              T tiling_data) {
    this->size = tiling_data.size;
    this->tileLength = tiling_data.tileLength;
    this->finalLength = tiling_data.finalLength;
    this->hasTanhx = tiling_data.hasTanhx;
    // 切分global
    GSetBuffer(GRAD, grad, 0, tiling_data.size);
    GSetBuffer(X, x, 0, tiling_data.size);
    if (hasTanhx) {
      GSetBuffer(TANHX, tanhx, 0, tiling_data.size);
    }
    GSetBuffer(Y, y, 0, tiling_data.size);
    // 初始化队列
    InitQueueSimple(GRAD, tileLength);
    InitQueueSimple(X, tileLength);
    if (hasTanhx) {
      InitQueueSimple(TANHX, tileLength);
    }
    InitQueueSimple(Y, tileLength);
    // 初始化BUF
    InitTBufBuffer(TEMP, tileLength * sizeof(float));
    if (!hasTanhx) {
      InitTBufBuffer(TEMP2, tileLength * sizeof(float));
      InitTBufBuffer(TEMP3, tileLength * sizeof(float));
    }
    if constexpr (std::is_same_v<_DT_X, half>) {
      InitTBufBuffer(CASTED_GRAD, tileLength * sizeof(float));
      InitTBufBuffer(CASTED_X, tileLength * sizeof(float));
      if (hasTanhx) {
        InitTBufBuffer(CASTED_TANHX, tileLength * sizeof(float));
      }
      InitTBufBuffer(CASTED_Y, tileLength * sizeof(float));
    }
  }

  __aicore__ inline void Process() {
    auto loopCount = size / tileLength;
    for (uint32_t i = 0; i < loopCount; ++i) {
      CopyIn(i, tileLength);
      Compute(tileLength);
      CopyOut(i, tileLength);
    }
    if (finalLength > 0) {
      auto copyCount = alignToBlock<_DT_X>(finalLength);
      CopyIn(loopCount, copyCount);
      Compute(finalLength);
      CopyOut(loopCount, copyCount);
    }
  }

  __aicore__ inline void CopyIn(uint32_t i, uint32_t calcCount) {
    EnQueGlobal2Local(GRAD, i * tileLength, calcCount);
    EnQueGlobal2Local(X, i * tileLength, calcCount);
    if (hasTanhx) {
      EnQueGlobal2Local(TANHX, i * tileLength, calcCount);
    }
  }

  __aicore__ inline void CopyOut(uint32_t i, uint32_t calcCount) {
    DeQueLocal2Global(Y, i * tileLength, calcCount);
  }

  __aicore__ inline void MishGradFloat(LocalTensor<float> &grad,
                                       LocalTensor<float> &x,
                                       LocalTensor<float> *tanhx,
                                       LocalTensor<float> &y,
                                       uint32_t calcCount) {
    if (hasTanhx) {
      TBufGet(TEMP, float);
      Mul(y, *tanhx, *tanhx, calcCount);
      Duplicate(BTensor(TEMP), float(1), alignToBlock<float>(calcCount));
      Sub(BTensor(TEMP), BTensor(TEMP), y, calcCount);
      Mul(BTensor(TEMP), BTensor(TEMP), x, calcCount);
      Exp(y, x, calcCount);
      Mul(BTensor(TEMP), BTensor(TEMP), y, calcCount);
      Adds(y, y, float(1), calcCount);
      Div(BTensor(TEMP), BTensor(TEMP), y, calcCount);
      Add(BTensor(TEMP), BTensor(TEMP), *tanhx, calcCount);
      Mul(y, BTensor(TEMP), grad, calcCount);
    } else {
      TBufGet(TEMP, float);
      TBufGet(TEMP2, float);
      TBufGet(TEMP3, float);
      // exp_x = np.exp(input_x)
      Exp(BTensor(TEMP), x, calcCount);
      // exp_add_x = exp_x + 1
      Adds(y, BTensor(TEMP), float(1), calcCount);
      // rec_exp_add_x = 1.0 / (exp_add_x * exp_add_x + 1.0)
      Mul(BTensor(TEMP2), y, y, calcCount);
      Adds(BTensor(TEMP2), BTensor(TEMP2), float(1), calcCount);
      Duplicate(BTensor(TEMP3), float(1), alignToBlock<float>(calcCount));
      Div(BTensor(TEMP2), BTensor(TEMP3), BTensor(TEMP2), calcCount);

      // result_1 = exp_add_x * exp_x * input_x * rec_exp_add_x * rec_exp_add_x
      // * 4
      Mul(BTensor(TEMP3), y, BTensor(TEMP), calcCount);
      Mul(BTensor(TEMP3), BTensor(TEMP3), x, calcCount);
      Mul(BTensor(TEMP3), BTensor(TEMP3), BTensor(TEMP2), calcCount);
      Mul(BTensor(TEMP3), BTensor(TEMP3), BTensor(TEMP2), calcCount);
      Muls(BTensor(TEMP3), BTensor(TEMP3), float(4), calcCount);

      // result = result_1 - 2 * rec_exp_add_x + 1
      Muls(BTensor(TEMP), BTensor(TEMP2), float(2), calcCount);
      Sub(BTensor(TEMP3), BTensor(TEMP3), BTensor(TEMP), calcCount);
      Adds(BTensor(TEMP3), BTensor(TEMP3), float(1), calcCount);
      // result = result * input_grad
      Mul(y, BTensor(TEMP3), grad, calcCount);
    }
  }

  __aicore__ inline void Compute(uint32_t calcCount) {
    if (hasTanhx) {
      DeQueSimple(GRAD);
      DeQueSimple(X);
      DeQueSimple(TANHX);
      QueAllocSimple(Y);
      if constexpr (std::is_same_v<_DT_GRAD, float>) {
        MishGradFloat(LTensor(GRAD), LTensor(X), &LTensor(TANHX), LTensor(Y),
                      calcCount);
      } else {
        TBufGet(CASTED_GRAD, float);
        TBufGet(CASTED_X, float);
        TBufGet(CASTED_TANHX, float);
        TBufGet(CASTED_Y, float);
        Cast(BTensor(CASTED_GRAD), LTensor(GRAD), RoundMode::CAST_NONE,
             calcCount);
        Cast(BTensor(CASTED_X), LTensor(X), RoundMode::CAST_NONE, calcCount);
        Cast(BTensor(CASTED_TANHX), LTensor(TANHX), RoundMode::CAST_NONE,
             calcCount);
        MishGradFloat(BTensor(CASTED_GRAD), BTensor(CASTED_X),
                      &BTensor(CASTED_TANHX), BTensor(CASTED_Y), calcCount);
        Cast(LTensor(Y), BTensor(CASTED_Y), RoundMode::CAST_NONE, calcCount);
      }
      EnQue(Y);
      QueFree(GRAD);
      QueFree(X);
      QueFree(TANHX);
    } else {
      DeQueSimple(GRAD);
      DeQueSimple(X);
      QueAllocSimple(Y);
      if constexpr (std::is_same_v<_DT_GRAD, float>) {
        MishGradFloat(LTensor(GRAD), LTensor(X), nullptr, LTensor(Y),
                      calcCount);
      } else {
        TBufGet(CASTED_GRAD, float);
        TBufGet(CASTED_X, float);
        TBufGet(CASTED_Y, float);
        Cast(BTensor(CASTED_GRAD), LTensor(GRAD), RoundMode::CAST_NONE,
             calcCount);
        Cast(BTensor(CASTED_X), LTensor(X), RoundMode::CAST_NONE, calcCount);
        MishGradFloat(BTensor(CASTED_GRAD), BTensor(CASTED_X), nullptr,
                      BTensor(CASTED_Y), calcCount);
        Cast(LTensor(Y), BTensor(CASTED_Y), RoundMode::CAST_NONE, calcCount);
      }
      EnQue(Y);
      QueFree(GRAD);
      QueFree(X);
    }
  }
};
} // namespace AscendC
extern "C" __global__ __aicore__ void mish_grad(GM_ADDR grad, GM_ADDR x,
                                                GM_ADDR tanhx, GM_ADDR y,
                                                GM_ADDR workspace,
                                                GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  AscendC::Kernel<DTYPE_GRAD, DTYPE_X, DTYPE_TANHX, DTYPE_Y> op;
  op.Init(grad, x, tanhx, y, tiling_data);
  op.Process();
}